from crowd_nav.configs.icra_benchmark.config import BaseEnvConfig, BasePolicyConfig, BaseTrainConfig, Config


class EnvConfig(BaseEnvConfig):
    def __init__(self, debug=False):
        super(EnvConfig, self).__init__(debug)


class PolicyConfig(BasePolicyConfig):
    def __init__(self, debug=False):
        super(PolicyConfig, self).__init__(debug)
        self.name = 'mp_stgcnn'

        # gcn
        self.gcn.num_layer = 2
        self.gcn.X_dim = 32

        self.gcn.similarity_function = 'embedded_gaussian'
        self.gcn.layerwise_graph = False
        self.gcn.skip_connection = True

        self.mp_stgcnn = Config()
        self.mp_stgcnn.linear_state_predictor = False
        self.mp_stgcnn.planning_depth = 1
        self.mp_stgcnn.planning_width = 1 # 10
        self.mp_stgcnn.do_action_clip = False
        self.mp_stgcnn.motion_predictor_dims = [64, 5]
        self.mp_stgcnn.value_network_dims = [32, 100, 100, 1]
        self.mp_stgcnn.share_graph_model = False

        self.mp_stgcnn.obs_seq_len = 4
        self.gcn.stgat_output_dim = 32  # predicted trajectry length * graph node output dimension


class TrainConfig(BaseTrainConfig):
    def __init__(self, debug=False):
        super(TrainConfig, self).__init__(debug)

        self.train.freeze_state_predictor = False
        self.train.detach_state_predictor = False
        self.train.reduce_sp_update_frequency = False